from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch as t

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.CenterCrop(50),
    transforms.RandomHorizontalFlip(),
    transforms.Normalize([0.485, 0.456, 0.5], [0.229, 0.224, 0.225])
])

def init_datasets(dataset_name, batch_size):
    if dataset_name == 'mnist':
        trainset = datasets.MNIST('./data/mnist', train=True, download=True, transform=transform)
        testset = datasets.MNIST('./data/mnist', train=False, download=True, transform=transform)
        trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)
        testloader = DataLoader(dataset=testset, batch_size=batch_size, shuffle=False)
        return trainloader, testloader
    if dataset_name == 'cifar10':
        trainset = datasets.CIFAR10('./data/cifar10', train=True, download=True, transform=transform)
        testset = datasets.CIFAR10('./data/cifar10', train=False, download=True, transform=transform)
        trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)
        testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)
        return trainloader, testloader
    if dataset_name == 'cifar100':
        trainset = datasets.CIFAR100('./data/cifar100', train=True, download=True, transform=transform)
        testset = datasets.CIFAR100('./data/cifar100', train=True, download=True, transform=transform)
        trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)
        testloader = DataLoader(dataset=testset, batch_size=batch_size, shuffle=False)
        return trainloader, testloader
if __name__ == '__main__':
    datasets_name = input()
    batch_size = int(input())
    init_datasets(datasets_name, batch_size)

